from typing import Generic, Type, TypeVar
import databases
from pydantic import BaseModel
from sqlalchemy import Select, func, select

from forward.extend.pydantic.base import BaseSchema

T = TypeVar("T", bound=BaseModel)


class PageSchema(BaseSchema, Generic[T]):
    page_num: int = 1
    page_size: int = 10
    total: int = 0
    lists: list[T] = []


class PageUtils:
    @staticmethod
    async def paginate(
        db: databases.Database,
        query: Select,
        schema: Type[T],
        page_num: int = 1,
        page_size: int = 10,
    ) -> PageSchema[T]:
        """
        @param db: 数据库连接
        @param query: 查询语句
        @param schema: 返回数据类型
        @param page_num: 页码
        @param page_size: 每页数量
        @return:
        """
        total = await db.fetch_val(
            select(func.count("*")).select_from(query.subquery())
        )
        records = await db.fetch_all(
            query.limit(page_size).offset((page_num - 1) * page_size)
        )
        lists = [
            schema.model_validate(record, from_attributes=True) for record in records
        ]
        return PageSchema[T](
            total=total, page_num=page_num, page_size=page_size, lists=lists
        )
